"""
Created by joe on 2019/12/25

voc xml 转 tensorflow_yolov3 txt
"""
import os
import sys

from bs4 import BeautifulSoup


class VocToTensorflowYolov3(object):
    """voc xml 转 tensorflow_yolov3 txt"""

    def __init__(self):
        """
        init
        """
        self.classes = []

    def transfer(self, label_path):
        """
        transfer
        :param label_path: 标签路径
        :return:
        """
        with open("{0}/{1}".format(label_path, "train_list.txt"), "w") as train_list_file:
            for filename in os.listdir(label_path):
                name_arr = filename.split(".")

                if not name_arr[1] == "xml":
                    continue

                xml_path = "{0}/{1}".format(label_path, filename)

                with open(xml_path, "r") as fi:
                    document = BeautifulSoup(fi.read(), "html.parser")

                    img_name = document.select("annotation filename")[0].text

                    train_list_file.write("{0}/{1}".format(label_path, img_name))

                    for obj_el in document.select("annotation object"):
                        class_name = obj_el.select("name")[0].text

                        if class_name not in self.classes:
                            self.classes.append(class_name)

                        x_min = int(obj_el.select("bndbox xmin")[0].text)
                        y_min = int(obj_el.select("bndbox ymin")[0].text)
                        x_max = int(obj_el.select("bndbox xmax")[0].text)
                        y_max = int(obj_el.select("bndbox ymax")[0].text)

                        train_list_file.write(
                            " {0},{1},{2},{3},{4}".format(x_min, y_min, x_max, y_max, self.classes.index(class_name))
                        )
                    train_list_file.write("\n")

        with open("{0}/classes.txt".format(label_path), "w") as fi:
            for class_name in self.classes:
                fi.write(class_name + "\n")


if __name__ == "__main__":
    """main"""
    transfer = VocToTensorflowYolov3()
    transfer.transfer(sys.argv[1])
